#include <Lucnx/stdlib.h>
#include <Lucnx/string.h>
#include <Lucnx/debug.h>
#include <Lucnx/types.h>
#include <Lucnx/debug.h>
#include <Lucnx/list.h>
#include <Lucnx/task.h>

#define NR_DESC 8

typedef struct block_t
{
    list_node node;
}block_t;

static block_desc_t kernel_desc_dir[NR_DESC]; // 内核的内存块描述符目录

void arena_init(block_desc_t *dir)
{
    u16 size = 16;
    for (size_t i = 0; i < NR_DESC; i++)
    {
        list_init(&dir[i].free_list);
        dir[i].size = size;
        size *= 2;
    }
}

void karena_init()
{
    arena_init(&kernel_desc_dir);
}

static block_t *block_ptr(int idx, arena_t *arena, u16 block_size)
{
    return (void*)((u32)&arena[1] + block_size * idx);
}

static arena_t *arena_ptr(void *block)
{
    return (arena_t*) ((u32)block & 0xFFFFF000);
}

extern void memory_lock_acquire(pool_flag pf);
extern void memory_lock_release(pool_flag pf);

void *kmalloc(u16 _Size)
{
    task_t *current = running_task();
    pool_flag pf = PF_KERNEL;
    memory_lock_acquire(pf); // 获取锁

    if (_Size >= PAGE_SIZE)
    {
        int nr = DIV_ROUND_UP(PAGE_SIZE + sizeof(arena_t), _Size);
        void *page = kpage_alloc(nr);
        arena_t *arena = page;
        arena->desc = NULL;
        arena->is_page = true;
        arena->nr = nr;
        memory_lock_release(pf);
        return &arena[1]; // 跨过 arena
    }

    block_desc_t *desc_dir = kernel_desc_dir;
    block_desc_t *b;

    // 选择合适大小的内存块
    for (size_t i = 0; i < NR_DESC; i++)
    {
        if (desc_dir[i].size > _Size)
        {
            b = &desc_dir[i]; 
            break;
        }
    }

    // 没有空闲内存块, 建立一个新的 arena
    if (list_null(&b->free_list))
    {
        arena_t *arena = kpage_alloc(1);
        memset(arena, 0, PAGE_SIZE); // 清空 arena
        int nr = (PAGE_SIZE - sizeof(arena_t)) / b->size; // 本 arena 可以存储多少内存块
        arena->nr = nr;
        arena->is_page = false; // 非整页分配
        arena->desc = b;        // 对应的内存块描述符

        // 加入空闲内存块链表
        for (size_t i = 0; i < nr; i++)
        {
            block_t *block = block_ptr(i, arena, b->size);
            list_add_end(&b->free_list, &block->node);
        }
    }
    list_node *block_node = list_pop(&b->free_list);
    block_t *mem_block = member2entry(block_t, node, block_node);

    memset(mem_block, 0, b->size);
    memory_lock_release(pf);
    return (void*)mem_block;
}

void kfree(void *address)
{
    pool_flag pf = PF_KERNEL;

    memory_lock_acquire(pf);

    arena_t *arena = arena_ptr(address);
    if (arena->is_page)
    {
        kpage_free(arena, arena->nr);
    }
    else
    {
        memset(address, 0, arena->desc->size);
        block_t *block = address;
        list_push(&arena->desc->free_list, &block->node);
    }
    memory_lock_release(pf);
}